rm(list = ls())

## 

library(rdrobust)
library(tidyverse)
library(pbapply)

## Declare covars

covs <- c('turnout_party_2009_btw', 'soz_vers_beschaeftigte_share',
          'pop_density_km2', 'migration_out_share')

## Get state incumbent

inc_lt <- read_rds("data/states_census.rds") %>% 
  dplyr::select(state_id, incumbent_pretreat_mp) %>% 
  dplyr::rename(inc_party_prior_to_census = incumbent_pretreat_mp)

## Get municipal election data

mu <- readRDS('data/data_municipal.rds') %>%
  filter(!is.na(treated)) %>%
  mutate(state_id = substr(ags, 1, 2)) %>% 
  filter(time_rel_period > -2)

## Relevant vars

# inc_share2 - incumbent share (only for the period right before and after census)
# inc_party_prior_to_census - incumbent prior to census

## RD w/ first differences - municipal 

olist <- c("agg_left", "agg_right", "cdu_csu",
           "spd",   "inc_share2")

## Source helper function to tidy rdrobust output

source("code/tidy_rd.R")

## Create first-differenced DF 

diff_df <- pblapply(olist, function(o) {
  out <- mu %>%
    filter(time_rel_period > -2) %>%
    mutate(time_rel_period = time_rel_period + 2)  %>% 
    pivot_wider(values_from = o, names_from = 'time_rel_period', 
                id_cols = 'ags',
                names_prefix = 'o') %>%
    mutate(diff = o2  - o1) %>%  dplyr::select(ags, diff) 
  ## Rename
  colnames(out)[2] <- o
  
  ## Return this
  out
}) %>%
  reduce(left_join) %>%
  left_join(mu %>% dplyr::select(ags, pop_dec_09, applies_census,
                                 one_of(covs), state_id,
                                 inc_party_prior_to_census) %>%
              distinct(ags, .keep_all = T)) %>%
  mutate(runvar = (pop_dec_09 * -1) + 10000) %>%
  filter(applies_census == 1)

## Create subsets

subset_select <- !diff_df$state_id == '08' 
subset_select_cdu <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'cdu_csu'
subset_select_spd <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'spd'

## Estimate across subsets

out_muni_full <- lapply(olist, function(o) {
  
  out_muni <- rdrobust(y = diff_df[subset_select, o] %>% pull(!!o), 
                       x = diff_df$runvar[subset_select], c = 0,
                       covs = diff_df[subset_select, covs])
  out_muni_cdu <- rdrobust(y = diff_df[subset_select_cdu, o] %>% pull(!!o), 
                           x = diff_df$runvar[subset_select_cdu], c = 0,
                           covs = diff_df[subset_select_cdu, covs])
  out_muni_spd <- rdrobust(y = diff_df[subset_select_spd, o] %>% pull(!!o), 
                           x = diff_df$runvar[subset_select_spd], c = 0,
                           covs = diff_df[subset_select_spd, covs])
  
  ## Done, state
  
  r1 <- out_muni %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'Full sample')
  r2 <- out_muni_cdu %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: CDU/CSU')
  r3 <- out_muni_spd %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: SPD')
  
  rbind(r1, r2, r3)
}) %>% 
  reduce(rbind) %>% 
  mutate(election = 'Municipal\nelections')

## Federal elections

bt <- readRDS("data/data_federal.rds") %>%
  filter(!is.na(treated)) %>% 
  filter(year > 2012)

## Outcomes

olist <- c("agg_left_party", "agg_right_party", "cdu_csu_party",
           "spd_party",
           "incumbent_party")

## First differences

diff_df <- pblapply(olist, function(o) {
  out <- bt %>%
    filter(year > 2012) %>%
    pivot_wider(values_from = o, names_from = 'year', id_cols = 'ags',
                names_prefix = 'o') %>%
    mutate(diff = o2017  - o2013) %>%  dplyr::select(ags, diff) 
  ## Rename
  colnames(out)[2] <- o
  
  ## Return this
  out
}) %>%
  reduce(left_join) %>%
  left_join(bt %>% dplyr::select(ags, pop_dec_09, applies_census,
                                 one_of(covs), state_id) %>%
              distinct(ags, .keep_all = T)) %>%
  mutate(runvar = (pop_dec_09 * -1) + 10000) %>%
  filter(applies_census == 1) %>% 
  filter(!state_id == '08')

## Merge state incumbent to this 

diff_df <- diff_df %>% 
  left_join(inc_lt)

## Subsets

subset_select <- !diff_df$state_id == '08' 
subset_select_cdu <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'cdu_csu'
subset_select_spd <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'spd'

## Estimate across subsets

out_btw_full <- lapply(olist, function(o) {
  
  
  out_btw <- rdrobust(y = diff_df[subset_select, o] %>% pull(!!o), 
                      x = diff_df$runvar[subset_select], c = 0,
                      covs = diff_df[subset_select, covs])
  out_btw_cdu <- rdrobust(y = diff_df[subset_select_cdu, o] %>% pull(!!o), 
                          x = diff_df$runvar[subset_select_cdu], c = 0,
                          covs = diff_df[subset_select_cdu, covs])
  out_btw_spd <- rdrobust(y = diff_df[subset_select_spd, o] %>% pull(!!o), 
                          x = diff_df$runvar[subset_select_spd], c = 0,
                          covs = diff_df[subset_select_spd, covs])
  
  ## Done, state
  
  r1 <- out_btw %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'Full sample')
  r2 <- out_btw_cdu %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: CDU/CSU')
  r3 <- out_btw_spd %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: SPD')
  
  rbind(r1, r2, r3)
}) %>% 
  reduce(rbind) %>% 
  mutate(election = 'Federal\nelections')

## State elections
## Get data

lt <- readRDS('data/data_state.rds') %>%
  filter(!is.na(treated)) %>%
  mutate(year = lubridate::year(date)) %>%
  mutate(state_id = substr(ags, 1, 2))

## Get info on states

lt <- read_rds("data/states_census.rds") %>%
  dplyr::select(applies_census, state_id, census_first_year) %>% 
  left_join(lt, .) %>%
  mutate(time_rel = year - census_first_year)

## Period rel. to treatment

periods_by_state <- lt %>%
  group_by(state_id) %>%
  summarise(l = list(unique(time_rel))) %>%
  unnest()
periods_by_state <- periods_by_state %>%
  group_by(state_id) %>%
  summarise(l_r = list(rank(l))) %>%
  unnest() %>% dplyr::select(-state_id) %>% 
  bind_cols(periods_by_state, .)
periods_by_state <- periods_by_state %>% 
  group_by(state_id) %>%
  mutate(l_r = l_r - max(l_r)) %>%
  filter(!is.na(l)) %>%
  rename(time_rel_period = l_r,
         time_rel = l)

## Merge this back to the main DF

lt <- lt %>%
  left_join(periods_by_state) %>%
  mutate(post = ifelse(time_rel_period > -1, 1, 0),
         treated_post = post * treated) %>% 
  filter(time_rel_period > -2)

## Declare outcomes

olist <- c("agg_left", "agg_right", 
           "cdu_csu",
           "spd",
           "incumbent")

## Get first differences

diff_df <- pblapply(olist, function(o) {
  out <- lt %>%
    filter(time_rel_period > -2) %>%
    mutate(time_rel_period = time_rel_period + 2)  %>% 
    pivot_wider(values_from = o, names_from = 'time_rel_period', 
                id_cols = 'ags',
                names_prefix = 'o') %>%
    mutate(diff = o2  - o1) %>%  dplyr::select(ags, diff) 
  ## Rename
  colnames(out)[2] <- o
  
  ## Return this
  out
}) %>%
  reduce(left_join) %>%
  left_join(lt %>% dplyr::select(ags, pop_dec_09, applies_census,
                                 one_of(covs),
                                 state_id) %>%
              distinct(ags, .keep_all = T)) %>%
  mutate(runvar = (pop_dec_09 * -1) + 10000) %>%
  filter(applies_census == 1) %>%
  filter(!str_sub(ags, 1, 2) == '01') %>% ## Exclude SH
  filter(!(substr(ags, 1, 2) == '08')) ## Exclude BW

## Merge state incumbent

diff_df <- diff_df %>% 
  left_join(inc_lt)

## Generate subsets

subset_select <- !diff_df$state_id == '08' 
subset_select_cdu <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'cdu_csu'
subset_select_spd <- !diff_df$state_id == '08' & 
  diff_df$inc_party_prior_to_census == 'spd'

## Estimate models

out_ltw_full <- lapply(olist, function(o) {
  
  out_ltw <- rdrobust(y = diff_df[subset_select, o] %>% pull(!!o), 
                      x = diff_df$runvar[subset_select], c = 0,
                      covs = diff_df[subset_select, covs])
  out_ltw_cdu <- rdrobust(y = diff_df[subset_select_cdu, o] %>% pull(!!o), 
                          x = diff_df$runvar[subset_select_cdu], c = 0,
                          covs = diff_df[subset_select_cdu, covs])
  out_ltw_spd <- rdrobust(y = diff_df[subset_select_spd, o] %>% pull(!!o), 
                          x = diff_df$runvar[subset_select_spd], c = 0,
                          covs = diff_df[subset_select_spd, covs])
  
  ## Done, state
  
  r1 <- out_ltw %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'Full sample')
  r2 <- out_ltw_cdu %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: CDU/CSU')
  r3 <- out_ltw_spd %>% tidy_rd(3) %>% 
    mutate(outcome = o, 
           incumbent_muni = 'State incumbent: SPD')
  
  rbind(r1, r2, r3)
}) %>% 
  reduce(rbind) %>% 
  mutate(election = 'State\nelections')

## Combine all results

r_full <- out_muni_full %>% 
  bind_rows(out_ltw_full) %>% 
  bind_rows(out_btw_full) %>% 
  mutate(outcome = str_remove_all(outcome, '_party')) %>% 
  mutate(outcome = ifelse(outcome == 'inc_share2', 'incumbent', outcome)) %>% 
  mutate(outcome = ifelse(outcome == 'incumbent' & str_detect(election, 'Federal'),
                          'fed. incumbent', outcome),
         outcome = ifelse(outcome == 'incumbent' & str_detect(election, 'State'), 
                          'state. incumbent',
                          outcome))

## Prep for plotting 

r_full <- r_full %>% 
  mutate(outcome= recode(outcome, 
                         `incumbent` = 'Municipal incumbent',
                         `state. incumbent` = 'State incumbent',
                         `fed. incumbent` = 'Federal incumbent',
                         `agg_left` = 'Left-wing parties',
                         `agg_center` = 'Centrist parties',
                         `agg_right` = 'Right-wing parties',
                         `spd` = 'SPD',
                         `cdu_csu` = 'CDU/CSU'))

## Aux DFs to order variables

order_df <- data.frame(outcome = c("Left-wing parties", "Right-wing parties", 
                                   "SPD", "CDU/CSU", 
                                   "Municipal incumbent",
                                   "State incumbent", 
                                   "Federal incumbent"),
                       order = rev(c(1, 2, 3, 4, 6, 7, 8)))

order_df_elec <- data.frame(election = c("Municipal\nelections", 
                                         "State\nelections", 
                                         "Federal\nelections"),
                            order_elec = rev(c(1, 2, 3)))

## Merge to the main results df

r_full <- r_full %>% 
  left_join(order_df) %>% 
  mutate(outcome = as.factor(outcome),
         outcome = fct_reorder(outcome, order)) %>% 
  left_join(order_df_elec) %>% 
  mutate(election = as.factor(election),
         election = fct_reorder(election, order_elec))

# Figure A.15: results conditional on state incumbent ----

r_full %>% 
  ggplot(aes(outcome, estimate)) +
  geom_hline(yintercept = 0, linetype = 'dotted') +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high),
                width = 0) +
  geom_point(shape = 21, fill = 'white') +
  facet_grid(election~incumbent_muni, scales = 'free_y') +
  theme_bw() +
  coord_flip() +
  ylab('RD estimate (percentage points)') +
  xlab('')

